import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import time
import requests
import zl_money_flow as zl
import stock_zt as zt


# ----------------------
# 1. 基础功能
# ----------------------
def is_trading_day(date):
    weekday = date.weekday()
    return weekday < 5

def get_current_period():
    now = datetime.now()
    today = now.date()
    if not is_trading_day(today):
        return "非交易日"
    morning_start = datetime(today.year, today.month, today.day, 9, 30)
    morning_end = datetime(today.year, today.month, today.day, 11, 30)
    afternoon_start = datetime(today.year, today.month, today.day, 13, 0)
    afternoon_end = datetime(today.year, today.month, today.day, 15, 0)
    if now < morning_start:
        return "盘前"
    elif morning_start <= now <= morning_end:
        return "上午交易中"
    elif morning_end < now < afternoon_start:
        return "午休（上午已收盘）"
    elif afternoon_start <= now <= afternoon_end:
        return "下午交易中"
    else:
        return "盘后"

def get_target_date():
    now = datetime.now()
    period = get_current_period()
    if period in ["盘前", "非交易日"]:
        target_date = (now - timedelta(days=1)).date()
        while not is_trading_day(target_date):
            target_date -= timedelta(days=1)
        return target_date
    else:
        return now.date()


# ----------------------
# 2. 数据获取
# ----------------------
def get_sina_latest_price(stock_code):
    if stock_code.startswith('6'):
        full_code = f"sh{stock_code}"
    else:
        full_code = f"sz{stock_code}"
    url = f"http://hq.sinajs.cn/list={full_code}"
    try:
        response = requests.get(url, timeout=5)
        if response.status_code == 200:
            data_str = response.text.split('=')[1].strip('";\n')
            data_list = data_str.split(',')
            if len(data_list) >= 4:
                current_price = float(data_list[3])
                prev_close = float(data_list[2])
                return {
                    "代码": stock_code,
                    "最新价格": current_price,
                    "昨日收盘价": prev_close,
                    "数据类型": "最新价"
                }
            else:
                print(f"{stock_code}数据字段不足（需至少4个，实际{len(data_list)}个）")
    except Exception as e:
        print(f"获取{stock_code}最新价格失败: {e}")
    return None


# ----------------------
# 3. 以涨停日为基准的指标计算与筛选
# ----------------------
def filter_stocks_by_zt_benchmark(hist_df):
    """以涨停日为基准，筛选资金留存高、股价表现差异大（资金股价差值大）的所有股票"""
    stock_metrics = []
    for code, group in hist_df.groupby("代码"):
        # 1. 筛选有涨停记录的股票（涨跌幅≥9.9%视为涨停）
        zt_candidates = group[group["涨跌幅"] >= 9.9]
        if zt_candidates.empty:
            continue  # 无涨停记录，排除
        zt_day = zt_candidates["日期"].min()  # 取最早涨停日作为基准日

        # 2. 获取涨停日数据（基准值）
        zt_day_data = group[group["日期"] == zt_day]
        if zt_day_data.empty:
            continue
        zt_close = zt_day_data["收盘价"].iloc[0]  # 涨停日收盘价（股价基准）
        zt_fund = zt_day_data["主力净流入-净额（万元）"].iloc[0]  # 涨停日主力资金（资金基准）
        if zt_fund <= 0:
            continue  # 涨停日资金未流入，排除

        # 3. 计算涨停日后的资金和股价指标（以涨停日为基准）
        group_sorted = group.sort_values("日期")
        # 仅保留涨停日之后的数据（分析涨停后的表现）
        post_zt_data = group_sorted[group_sorted["日期"] >= zt_day]
        if len(post_zt_data) < 2:
            continue  # 涨停后无足够数据，排除

        # 计算累计资金（以涨停日为基准100%）
        group_sorted["累计主力资金"] = group_sorted["主力净流入-净额（万元）"].cumsum()
        zt_cum_fund = group_sorted[group_sorted["日期"] == zt_day]["累计主力资金"].iloc[0]
        group_sorted["资金留存率(%)"] = (group_sorted["累计主力资金"] / zt_cum_fund) * 100  # 资金留存比例

        # 计算股价相对涨停日的涨幅（修改：从跌幅改为涨幅，保持正负逻辑）
        group_sorted["股价相对涨停日(%)"] = (group_sorted["收盘价"] / zt_close) * 100  # 现价相对涨停日的比例
        group_sorted["股价涨幅(%)"] = group_sorted["股价相对涨停日(%)"] - 100  # 涨幅（正值为上涨，负值为下跌）

        # 4. 取最近一个交易日的数据（最新表现）
        latest_data = group_sorted.iloc[-1]
        if latest_data["日期"] == zt_day:
            continue  # 仍为涨停日，无后续数据，排除

        # 核心指标：资金留存率 - 股价相对比例（差值越大越好）
        fund_price_diff = latest_data["资金留存率(%)"] - latest_data["股价相对涨停日(%)"]

        # 5. 保存指标
        stock_metrics.append({
            "代码": code,
            "名称": group["名称"].iloc[0],
            "涨停日": zt_day.strftime('%Y-%m-%d'),
            "资金留存率(%)": latest_data["资金留存率(%)"],  # 越高越好
            "股价涨幅(%)": latest_data["股价涨幅(%)"],      # 可正可负（正值为上涨，负值为下跌）
            "资金股价差值(%)": fund_price_diff  # 核心筛选指标（越大越好）
        })

    # 按“资金股价差值”排序
    if not stock_metrics:
        return []
    sorted_stocks = sorted(stock_metrics, key=lambda x: x["资金股价差值(%)"], reverse=True)
    return sorted_stocks


# ----------------------
# 4. 主逻辑（整合筛选+最新价格获取）
# ----------------------
def get_realtime_or_latest_data(codes, hist_df, name_map):
    period = get_current_period()
    target_date = get_target_date()
    print(f"\n当前时段：{period} → 目标数据日期：{target_date}")

    if period in ["上午交易中", "午休（上午已收盘）", "下午交易中"]:
        print(f"获取{period}的实时数据（返回最新价格）...")
        latest_data = []
        for code in codes:
            stock_data = get_sina_latest_price(code)
            if stock_data:
                if stock_data["最新价格"] <= 0:
                    print(f"{code}价格异常，跳过")
                    continue
                latest_data.append({
                    "代码": code,
                    "当前价格": stock_data["最新价格"],
                    "基准价格": stock_data["昨日收盘价"],
                    "数据类型": stock_data["数据类型"],
                    "数据日期": target_date
                })
                print(f"{code}（{name_map[code]}）最新价格：{stock_data['最新价格']}")
            else:
                print(f"{code}（{name_map[code]}）未获取到最新价格，跳过")
            time.sleep(0.3)
        return target_date, latest_data

    else:
        print(f"获取{target_date}收盘数据...")
        close_data = []
        for code in codes:
            code_hist = hist_df[hist_df["代码"] == code]
            if code_hist.empty:
                print(f"{code}无历史数据，跳过")
                continue
            target_data = code_hist[code_hist["日期"] == target_date]
            if target_data.empty:
                target_data = code_hist[code_hist["日期"] == max(code_hist["日期"])]
                print(f"{code}今日数据未更新，使用最近历史数据")
            if not target_data.empty:
                close_price = target_data["收盘价"].iloc[0]
                prev_dates = code_hist[code_hist["日期"] < target_date]["日期"]
                prev_close = code_hist[code_hist["日期"] == max(prev_dates)]["收盘价"].iloc[
                    0] if not prev_dates.empty else close_price
                close_data.append({
                    "代码": code,
                    "当前价格": close_price,
                    "基准价格": prev_close,
                    "数据类型": "最新价",
                    "数据日期": target_date
                })
        return target_date, close_data


# ----------------------
# 5. 执行与输出（完整流程）
# ----------------------
if __name__ == "__main__":
    try:
        # 1. 加载历史数据
        hist_df = zl.all_past_5_day_data
        hist_df["代码"] = hist_df["代码"].apply(lambda x: str(x).zfill(6))
        hist_df["日期"] = pd.to_datetime(hist_df["日期"]).dt.date
        print(f"成功加载历史数据（共{len(hist_df)}条记录，包含{hist_df['代码'].nunique()}只股票）")

        # 2. 核心：筛选“涨停后资金留存高、股价表现差异大”的所有股票
        filtered_stocks = filter_stocks_by_zt_benchmark(hist_df)
        if not filtered_stocks:
            print("无符合条件的股票（需有涨停记录，且涨停后有资金留存）")
            exit()
        print(f"\n已筛选出{len(filtered_stocks)}支符合条件的股票（资金留存高、股价表现差异大）：")
        for i, stock in enumerate(filtered_stocks[:5]):  # 仍只打印前5支，避免输出过长
            price_change = stock["股价涨幅(%)"]
            price_change_str = f"{'上涨' if price_change >= 0 else '下跌'}{abs(price_change):.2f}%"
            print(f"  排名{i+1}：{stock['代码']}（{stock['名称']}）→ "
                  f"资金留存率{stock['资金留存率(%)']:.2f}%，股价{price_change_str}，"
                  f"资金股价差值{stock['资金股价差值(%)']:.2f}%")
        if len(filtered_stocks) > 5:
            print(f"  ...（省略{len(filtered_stocks)-5}支）")

        # 3. 提取筛选后的股票代码和名称
        all_codes = [stock["代码"] for stock in filtered_stocks]
        name_map = {stock["代码"]: stock["名称"] for stock in filtered_stocks}

        # 4. 获取所有符合条件股票的最新价格
        target_date, target_data = get_realtime_or_latest_data(all_codes, hist_df, name_map)

        # 5. 整理结果（整合筛选指标和最新价格）
        result_list = []
        for item in target_data:
            code = item["代码"]
            current_price = item["当前价格"]
            base_price = item["基准价格"]
            pct_change = (current_price - base_price) / base_price * 100  # 当日涨幅
            # 补充筛选时的核心指标
            stock_info = next(s for s in filtered_stocks if s["代码"] == code)
            result_list.append({
                "股票代码": code,
                "股票名称": name_map[code],
                "涨停日": stock_info["涨停日"],
                "资金留存率(%)": round(stock_info["资金留存率(%)"], 2),
                "股价涨幅(%)": round(stock_info["股价涨幅(%)"], 2),  # 修改：从跌幅改为涨幅
                "昨日收盘价(元)": round(base_price, 2),
                "最新价格(元)": current_price,
                "当日涨幅(%)": round(pct_change, 2),
                "记录时间": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            })

        # 6. 保存并输出
        if result_list:
            result_df = pd.DataFrame(result_list)
            result_df.to_csv(zt.get_file_path("结果排名.csv"), index=False, encoding="utf-8-sig")
            print(f"\n✅ 已将{len(result_list)}支股票数据保存至：{zt.get_file_path('结果排名.csv')}")
            print("\n===== 最终结果（前10支） =====")  # 仅显示前10支，避免控制台输出过长
            print(result_df.sort_values("资金留存率(%)", ascending=False).head(10)[
                ["股票代码", "股票名称", "资金留存率(%)", "股价涨幅(%)", "最新价格(元)", "当日涨幅(%)"]
            ])
            print(f"\n完整结果已保存至CSV文件（共{len(result_list)}支股票）")
        else:
            print("无有效最新价格数据")

    except Exception as e:
        print(f"程序执行失败：{e}")